#!/usr/bin/python3
# -*- coding: utf-8 -*-

import time
import uuid

class SnowflakeID:
    """
    雪花算法ID生成器
    
    基于Twitter的Snowflake算法实现的分布式全局唯一ID生成器。
    ID结构：1位符号位 + 41位时间戳 + 5位数据中心ID + 5位工作机器ID + 12位序列号
    """
    
    # 开始时间戳 (2020-01-01 00:00:00 UTC)
    EPOCH = 1577836800000
    
    # 位长度
    WORKER_ID_BITS = 5
    DATACENTER_ID_BITS = 5
    SEQUENCE_BITS = 12
    
    # 最大值
    MAX_WORKER_ID = -1 ^ (-1 << WORKER_ID_BITS)  # 31
    MAX_DATACENTER_ID = -1 ^ (-1 << DATACENTER_ID_BITS)  # 31
    MAX_SEQUENCE = -1 ^ (-1 << SEQUENCE_BITS)  # 4095
    
    # 位移
    WORKER_ID_SHIFT = SEQUENCE_BITS  # 12
    DATACENTER_ID_SHIFT = SEQUENCE_BITS + WORKER_ID_BITS  # 17
    TIMESTAMP_SHIFT = SEQUENCE_BITS + WORKER_ID_BITS + DATACENTER_ID_BITS  # 22
    
    def __init__(self, worker_id=0, datacenter_id=None):
        """
        初始化雪花ID生成器
        
        Args:
            worker_id: 工作机器ID (0-31)
            datacenter_id: 数据中心ID (0-31)，默认与worker_id相同
        """
        # 验证worker_id
        if worker_id > self.MAX_WORKER_ID or worker_id < 0:
            raise ValueError(f"Worker ID不能大于{self.MAX_WORKER_ID}或小于0")
        self.worker_id = worker_id
        
        # 验证datacenter_id
        if datacenter_id is None:
            datacenter_id = worker_id
        if datacenter_id > self.MAX_DATACENTER_ID or datacenter_id < 0:
            raise ValueError(f"Datacenter ID不能大于{self.MAX_DATACENTER_ID}或小于0")
        self.datacenter_id = datacenter_id
        
        # 序列号
        self.sequence = 0
        # 上次生成ID的时间戳
        self.last_timestamp = -1
    
    def _generate_timestamp(self):
        """获取当前时间戳（毫秒）"""
        return int(time.time() * 1000)
    
    def _wait_next_millis(self, last_timestamp):
        """等待到下一毫秒"""
        timestamp = self._generate_timestamp()
        while timestamp <= last_timestamp:
            timestamp = self._generate_timestamp()
        return timestamp
    
    def next_id(self):
        """
        生成下一个ID
        
        Returns:
            生成的唯一ID（整数）
        """
        timestamp = self._generate_timestamp()
        
        # 如果当前时间小于上一次ID生成的时间，说明系统时钟回退，抛出异常
        if timestamp < self.last_timestamp:
            raise RuntimeError(f"时钟回退，拒绝生成ID，上次时间: {self.last_timestamp}, 当前时间: {timestamp}")
        
        # 如果是同一时间生成的，则进行序列号递增
        if timestamp == self.last_timestamp:
            # 序列号递增
            self.sequence = (self.sequence + 1) & self.MAX_SEQUENCE
            # 序列号溢出
            if self.sequence == 0:
                # 阻塞到下一个毫秒
                timestamp = self._wait_next_millis(self.last_timestamp)
        else:
            # 不同毫秒内，序列号重置为0
            self.sequence = 0
        
        # 更新上次生成ID的时间戳
        self.last_timestamp = timestamp
        
        # 组合ID
        return ((timestamp - self.EPOCH) << self.TIMESTAMP_SHIFT) | \
               (self.datacenter_id << self.DATACENTER_ID_SHIFT) | \
               (self.worker_id << self.WORKER_ID_SHIFT) | \
               self.sequence
    
    def format_id(self):
        """
        生成带有Freyr前缀的格式化ID
        
        Returns:
            str: 格式化的ID，格式为 'FREYR-{雪花ID}'
        """
        return f"FREYR-{self.next_id()}"

class UuidGenerator:
    """
    UUID生成器
    
    基于Python标准库uuid模块的UUID生成器。
    默认使用UUID4 (随机生成)。
    """
    
    def __init__(self, prefix=None):
        """
        初始化UUID生成器
        
        Args:
            prefix: 可选的ID前缀
        """
        self.prefix = prefix
    
    def next_id(self):
        """
        生成下一个UUID
        
        Returns:
            str: 生成的UUID字符串
        """
        return str(uuid.uuid4())
    
    def next_id_no_hyphen(self):
        """
        生成下一个不包含短横线的UUID
        
        Returns:
            str: 不包含短横线的UUID字符串
        """
        return str(uuid.uuid4()).replace('-', '')
    
    def format_id(self):
        """
        生成带有可选前缀的格式化UUID
        
        Returns:
            str: 格式化的UUID，如果有前缀则格式为 '{prefix}-{uuid}'，否则直接返回uuid
        """
        uuid_str = self.next_id()
        if self.prefix:
            return f"{self.prefix}-{uuid_str}"
        return uuid_str
    
    def format_id_no_hyphen(self):
        """
        生成不包含短横线的UUID，可选带前缀
        
        Returns:
            str: 不包含短横线的UUID，如果有前缀则格式为 '{prefix}{uuid_no_hyphen}'
        """
        uuid_str = self.next_id_no_hyphen()
        if self.prefix:
            return f"{self.prefix}{uuid_str}"
        return uuid_str

# 创建ID生成器单例
class IdGeneratorSingleton:
    """ID生成器单例类"""
    _snowflake_instance = None
    _uuid_instance = None
    _generator_type = "snowflake"  # 默认使用雪花算法
    
    @classmethod
    def set_generator_type(cls, generator_type):
        """设置ID生成器类型
        
        Args:
            generator_type: 生成器类型，支持 'snowflake' 或 'uuid'
        """
        if generator_type not in ['snowflake', 'uuid']:
            raise ValueError("生成器类型必须是 'snowflake' 或 'uuid'")
        cls._generator_type = generator_type
    
    @classmethod
    def get_instance(cls, worker_id=1, prefix=None):
        """获取ID生成器实例（单例模式）
        
        Args:
            worker_id: 雪花算法的工作机器ID
            prefix: UUID的可选前缀
            
        Returns:
            SnowflakeID or UuidGenerator: 根据当前设置返回相应的生成器实例
        """
        if cls._generator_type == 'snowflake':
            if cls._snowflake_instance is None:
                cls._snowflake_instance = SnowflakeID(worker_id=worker_id)
            return cls._snowflake_instance
        else:
            if cls._uuid_instance is None:
                cls._uuid_instance = UuidGenerator(prefix=prefix)
            return cls._uuid_instance 